
import sys
import indexing.featureindex
import coverage.readsperbasefromsummary
import coverage.basecountreader

def extractCountsFromSortedSummary(locfile, summaryfile, sample):
    ix = indexing.featureindex.FeatureIndex(summaryfile)

    for line in open(locfile):
        cols = line.strip().split("\t")
        (chr, pos) = cols[0].split(":")
        if "-" in pos:
            (minpos, maxpos) = pos.split("-")
            minpos = long(minpos)
            maxpos = long(maxpos)
        else:
            minpos = maxpos = long(pos)
        features = list(coverage.readsperbasefromsummary.process((chr, '?', minpos, maxpos), ix))
        res = ["" for p in xrange(minpos, maxpos+1)]
        for (p, feature) in features:
            res[p-minpos] = "%s=%s" % (sample, "%d/%d/%d/%d" % feature.getFragmentCounts())
        for p in xrange(minpos, maxpos+1):
            print "%s:%ld\t%s" % (chr, p, "\t".join(cols[1:] + [res[p-minpos]]))

def extractSnpsFromSortedSummary(locfile, summaryfile, minVariantFrags=2):
    ix = indexing.featureindex.FeatureIndex(summaryfile)
    n = 0
    for line in open(locfile):
        cols = line.strip().split("\t")
        (chr, pos) = cols[0].split(":")
        if "-" in pos:
            (minpos, maxpos) = pos.split("-")
            minpos = long(minpos)
            maxpos = long(maxpos)
        else:
            minpos = maxpos = long(pos)
        n += 1
        # if 0==(n%10): print >> sys.stderr, "# Query", n, "=", (chr, '?', minpos, maxpos)
        coverage.readsperbasefromsummary.printAsSnps((chr, '?', minpos, maxpos), ix, minVariantFrags)
            
def extractCountsStreaming(locfile, covfile, sample):
    currChr = None
    ix = indexing.featureindex.FeatureIndex(covfile)
    reader = coverage.basecountreader.BaseCountReader(ix)
    for line in open(locfile):
        cols = line.strip().split("\t")
        (chr, pos) = cols[0].split(":")
        pos = long(pos)
        feature = reader[(chr, pos)]
        if not feature:
            print "%s\t%s" % (line.strip(), "")
        else:
            print "%s\t%s=%s" % (line.strip(), sample, "%d/%d/%d/%d" % feature.getDetails().getFragmentCounts())

if __name__=="__main__":
    asSnps = False
    minVariantFrags = 2
    if sys.argv[1] == "--as-snps":
        asSnps = True
        sys.argv = sys.argv[1:]
    if sys.argv[1].startswith("--min-frags="):
        minVariantFrags = int(sys.argv[1].split("=")[-1])
        print >> sys.stderr, "# Using minimum variant fragments = ", minVariantFrags
        sys.argv = sys.argv[1:]
    locfile = sys.argv[1]
    summaryfile = sys.argv[2]
    if asSnps:
        extractSnpsFromSortedSummary(locfile, summaryfile, minVariantFrags)
    else:
        sample = sys.argv[3]
        extractCountsFromSortedSummary(locfile, summaryfile, sample)
    # extractCountsStreaming(locfile, summaryfile, sample)
